# -*- coding: utf-8 -*-
"""
该代码在removeDup之后执行

1. 将数据库中的新闻进行分词
2. 在新闻中寻找匹配的股票，对新闻进行打标签
3. 将打完标签后的新闻重新写回数据库

author：yinzm
"""
import pymongo
import jieba
from string import punctuation
import re
import pickle
import os
# 载入自定义词典
jieba.load_userdict("../ExtractWindInfo/dictionary/biggerFinancialDict.csv")

# 所有的标点符号
punc = punctuation + '.,;《》？！“”‘’@#￥%…&×（）—+【】{};；●，。·&～、|\:：'

# 建立MongoDB连接
mongoClient = pymongo.MongoClient("localhost", 27017)
db = mongoClient.sinafinance
collection = db.sinafinance2

newCollection = db.sinafinance3 # 将数据转存到该数据库中

# 将所有股票与实体之间的映射关系加载进来
pklPath = '../ExtractWindInfo/data/pklFiles'

# 按照日期将所有的映射关系加载进来
date_wind = []
for year in range(2012, 2018):
    for date_str in ['0501', '0901', '1101']:
        date_wind.append(str(year)+date_str)

relation = {} # relation['20120501'] = {'boardchairman':boardchairman, 'ceo':ceo, ....}
for date in date_wind:
    filePath = os.path.join(pklPath, date)
    relation[date] = {}
    boardchairman = pickle.load(open(os.path.join(filePath, 'boardchairman.pkl'), 'rb'))
    relation[date]['boardchairman'] = boardchairman
    ceo = pickle.load(open(os.path.join(filePath, 'ceo.pkl'), 'rb'))
    relation[date]['ceo'] = ceo
    chairman = pickle.load(open(os.path.join(filePath, 'chairman.pkl'), 'rb'))
    relation[date]['chairman'] = chairman
    compName = pickle.load(open(os.path.join(filePath, 'compName.pkl'), 'rb'))
    relation[date]['compName'] = compName
    holderController = pickle.load(open(os.path.join(filePath, 'holderController.pkl'), 'rb'))
    relation[date]['holderController'] = holderController
    holderLiqname = pickle.load(open(os.path.join(filePath, 'holderLiqname.pkl'), 'rb'))
    relation[date]['holderLiqname'] = holderLiqname
    holderName = pickle.load(open(os.path.join(filePath, 'holderName.pkl'), 'rb'))
    relation[date]['holderName'] = holderName
    preName = pickle.load(open(os.path.join(filePath, 'preName.pkl'), 'rb'))
    relation[date]['preName'] = preName
    secName = pickle.load(open(os.path.join(filePath, 'secName.pkl'), 'rb'))
    relation[date]['secName'] = secName
    secName1 = pickle.load(open(os.path.join(filePath, 'secName1.pkl'), 'rb'))
    relation[date]['secName1'] = secName1
    stockCodeDiffVariation = pickle.load(open(os.path.join(filePath, 'stockCodeDiffVariation.pkl'), 'rb'))
    relation[date]['stockCodeDiffVariation'] = stockCodeDiffVariation


# 加载所有的股票列表
stockCodeList = []
with open('../ExtractWindInfo/data/stockList.csv', 'r', encoding='utf-8') as f:
    for stock in f.readlines():
        stockCode = stock.strip()[:-3]
        stockCodeList.append(stockCode)


# 从数据库中抽取数据，然后开始处理
cnt = 1
for item in collection.find():
    new = item['content']
    news_date = item['date']
    # 找到该新闻所能用到的最新的relation
    relation_date = None
    for i in range(1, len(date_wind)):
        if news_date > date_wind[i-1] and news_date <= date_wind[i]:
            relation_date = date_wind[i-1]
            break
    if relation_date == None:
        continue

    new2 = re.sub('[{}]'.format(punc), '', new) # 去除所有的空格还有标点符号
    out = list(jieba.cut(new2))# 这个用来在pairMap中查找
    item['content'] = "/".join(out) # 这个存到数据库中去

    # 开始查找新闻影响的股票
    stock_tag = set() # 这条新闻所影响的股票集合
    for word in out:
        if word in stockCodeList:
            stock_tag.add(word)
        elif word in relation[relation_date]['boardchairman']:
            stock_tag.add(relation[relation_date]['boardchairman'][word])
        elif word in relation[relation_date]['ceo']:
            stock_tag.add(relation[relation_date]['ceo'][word])
        elif word in relation[relation_date]['chairman']:
            stock_tag.add(relation[relation_date]['chairman'][word])
        elif word in relation[relation_date]['compName']:
            stock_tag.add(relation[relation_date]['compName'][word])
        elif word in relation[relation_date]['holderController']:
            stock_tag.add(relation[relation_date]['holderController'][word])
        elif word in relation[relation_date]['holderLiqname']:
            stock_tag.add(relation[relation_date]['holderLiqname'][word])
        elif word in relation[relation_date]['holderName']:
            stock_tag.add(relation[relation_date]['holderName'][word])
        elif word in relation[relation_date]['preName']:
            stock_tag.add(relation[relation_date]['preName'][word])
        elif word in relation[relation_date]['secName']:
            stock_tag.add(relation[relation_date]['secName'][word])
        elif word in relation[relation_date]['secName1']:
            stock_tag.add(relation[relation_date]['secName1'][word])
        elif word in relation[relation_date]['stockCodeDiffVariation']:
            stock_tag.add(relation[relation_date]['stockCodeDiffVariation'][word])

    item['stock'] = ','.join(list(stock_tag))
    # 写回到新的数据库
    newCollection.insert(item)
    # 打印进度
    if cnt%100 == 0:
        print(cnt)
    cnt += 1



